"""符号处理工具函数."""

from __future__ import annotations

import re
from typing import Optional


_EXCHANGE_SUFFIXES = {
    "sz": ".SZ",
    "sh": ".SS",
    "ss": ".SS",
    "hk": ".HK",
    "us": "",
}


def normalize_symbol(symbol: str) -> str:
    """将 ETF 代码转换为 AkShare 可识别的格式（6位数字代码）.

    Examples
    --------
    >>> normalize_symbol("159919")
    '159919'
    >>> normalize_symbol("515220")
    '515220'
    >>> normalize_symbol("sz159919")
    '159919'
    >>> normalize_symbol("sh510300")
    '510300'
    >>> normalize_symbol("515790.SS")
    '515790'
    >>> normalize_symbol("spy")
    'SPY'
    """

    if not symbol:
        raise ValueError("symbol 不能为空")

    symbol = symbol.strip()
    
    # 如果包含点号，去掉后缀（如 515790.SS -> 515790）
    if "." in symbol:
        symbol = symbol.split(".")[0]

    lowered = symbol.lower()
    match = re.match(r"^(sz|sh|ss|hk|us)?(\d{6}|\w+)$", lowered)
    if not match:
        return symbol.upper()

    prefix, body = match.groups()
    
    # 返回6位数字代码（AkShare 格式）
    # 如果是非数字代码（如美股），返回大写
    if body.isdigit():
        return body.upper()
    else:
        return symbol.upper()


def extract_exchange(symbol: str) -> Optional[str]:
    """根据符号推断交易所."""

    if symbol.endswith(".SZ"):
        return "SZSE"
    if symbol.endswith(".SS"):
        return "SSE"
    if symbol.endswith(".HK"):
        return "HKEX"
    if symbol.endswith(".L"):
        return "LSE"
    if symbol.endswith(".T"):
        return "TSE"
    if symbol.isupper() and len(symbol) <= 5:
        return "NYSE/NASDAQ"
    return None

